#!/usr/bin/python

## TODO

import os.path                  # to manipulate filenames
import sys                      # error and output messages
import re                       # checking date formats
import numpy                    # needed to hold array for output and do maths
import numpy.ma                 # needed to hold array for output and do maths
import getopt                   # CLI options

from tcclimFunctions import *   # get common functions 

def getCruTS2(variable, start='', end='', outbase='', datefile='', average=True, verbose=True, format='ascii', transposedb=''):
    
    
    # variable = 'pre'
    # start = '01,1951'
    # end = '01,1952'
    # datefile=''; average=True; verbose=True;
    
    # This is the demo example given in the guidance notes
    # getTynSC2('pre', '12,2002 ', '12,2002 ', average=False)
    # test point (363, 286) should be 90.0
    
    
    #########################
    # SETUP AND TEST INPUTS #
    #########################
    
    variables = ['dtr','pre','tmp','vap','cld','tmx','tmn','frs','wet']
    
    if not variable in variables:
        sys.stderr.write('Variable \'' + variable + '\' not recognised\n')
        sys.exit()
    
    if not os.path.exists(transposedb):
        sys.stderr.write('Transpose database not found at: ' + os.path.abspath(os.path.join(os.getcwd(), transposedb)) + '\n')
        sys.exit()
    
    if verbose: 
        sys.stdout.write('Obtaining dates...\n')
    
    ##########################
    # GET THE RANGE OF DATES #
    ##########################
    
    # compile a regular expression to check for badly formatted count fields
    reMoYr = re.compile('[0-9]{1,2},[0-9]{4}') 
    
    # turn start and end into a list or read from file
    if datefile != '':
        # try and get the file contents
        try:
            dateFile = open(datefile, 'r')
        except IOError:
            sys.stderr.write('Cannot open file containing dates: ' + datefile + '\n')
            sys.exit()
        
        dateKeys = dateFile.readlines()
        dateKeys = [m.strip() for m in dateKeys]
        dateFile.close()
        
        # test for carriage return at end
        if dateKeys[-1] == '':
            dateKeys = dateKeys[:-1]
        
        # test for two parts
        dates = [m.split(',') for m in dateKeys]
        nPart = [len(m) == 2 for m in dates]
        if not all(nPart):
            sys.stderr.write('Badly formatted dates in: ' + datefile + '\n')
            sys.exit()
        
        # try and convert pairs of values to integers
        try:
            dates = [(int(m[0]), int(m[1])) for m in dates]
        except:
            sys.stderr.write('Datefile contains non-integer values: ' + datefile + '\n')
            sys.exit()
        
        # check their range
        monthValid = [m[0] > 0 and m[0] <= 12  for m in dates]
        yearValid = [m[1] > 1900 and m[1] <= 2002  for m in dates]
        
        if not all(monthValid):
            sys.stderr.write('Month in ' + datefile + ' outside range [1-12].\n')
            sys.exit()
        if not all(yearValid):
            sys.stderr.write('Year in ' + datefile + ' outside range [1901,2002].\n')
            sys.exit()
        
        # remove duplicates
        noDupes = list(set(dates))
        if len(noDupes) < len(dates):
            sys.stderr.write('WARNING: duplicate dates removed from dates in ' + datefile + '\n')
            dates = noDupes
            dates.sort()
        
    else: 
        # are both start and end defined and defined correctly
        if start == '' or end == '':
            sys.stderr.write('Missing start or end date.\n')
            sys.exit()
        
        if not reMoYr.match(start):
            sys.stderr.write('Misformatted start date: ' + start + '\n')
            sys.exit()
        
        if not reMoYr.match(end):
            sys.stderr.write('Misformatted end date: ' + end + '\n')
            sys.exit()
        
        # check the dates make sense
        start = map(int, start.split(','))
        end   = map(int, end.split(','))
        if start[0] > 12:
            sys.stderr.write('Bad start month: ' + str(start[0]) + '\n')
            sys.exit()
        
        if end[0] > 12:
            sys.stderr.write('Bad end month: ' + str(end[0]) + '\n')
            sys.exit()
        
        if start[1] < 1901 or start[1] > 2002:
            sys.stderr.write('Start year not in range 1901 - 2002: ' + str(start[1]) + '\n')
            sys.exit()
        
        if end[1] < 1901 or end[1] > 2002:
            sys.stderr.write('End year not in range 1901 - 2002: ' + str(end[1]) + '\n')
            sys.exit()
        
        # get the date range
        if (start[1] > end[1]) or (start[1] == end[1] and start[0] > end[0]):
            sys.stderr.write('Start date later than end date\n')
            sys.exit()
        
        start_month = start[0]
        nMonths     = (end[1] - start[1]) * 12 + end[0] + 1
        dates = [((m - 1) % 12 + 1, (m - 1) / 12 + start[1]) for m in range(start_month, nMonths)]
    
    # reconsitute dateKeys from parsed dates - has the effect of stripping leading zeros
    dateKeys = [str(m[0]) + ',' + str(m[1]) for m in dates]
    
    ########################
    # READ DATA FROM FILES #
    ########################
    
    # read variable file:
    # - monthly maps of observed values (1901-2002) for variables
    # - dated [month],1990
    # - this is time consuming because the files are so large
    #   but this now only returns the dates needed
    if verbose: sys.stdout.write('Reading variable values for ' + str(len(dateKeys)) + ' months...\n')
    
    varFile = os.path.join(transposedb, 'tr.cru_ts_2_10.1901-2002.' + variable)
    transData = readTransposeGrim(varFile, dateKeys, ticker=verbose)
    
    # STORE DATA IN A NUMPY ARRAY
    # - this _does_ create 2d array with only 1 date
    valArray = numpy.zeros([len(dates), 67420])
    
    for currDate in range(len(dates)):
        c_key   = dateKeys[currDate]
        valArray[currDate,] = transData[c_key]
    
    ## SCALE THE DATA (NB - no need to check for clipping in this data - no projection)
    scaleVars = {'tmp': 0.1, 'pre': 0.1, 'dtr': 0.1, 'vap': 0.1, 'cld': 0.1, 'tmn':0.1, 'tmx':0.1, 'frs':0.01, 'wet':0.01}
    valArray = valArray * scaleVars[variable]
    
    # average across the dates if required (axis 0 = columns)
    #  - returns a masked value if _all_ values in a column are masked
    #    but otherwise returns a non-masked array, so need to remask
    if average:
        if verbose: sys.stdout.write('Averaging values across dates...\n')
        valArray = numpy.average(valArray, axis=0)
    
    # handle singleton year extractions and averages, where numpy
    # drops dimensions, by putting the dimension back in
    if len(valArray.shape) == 1:
        valArray = valArray[numpy.newaxis, :]
    
    ##########################
    # WRITE OUTPUTS TO FILES #
    ##########################
    
    # get the base filename
    if outbase == '':
        outbase = variable  + '.'
    
    # now output the data as ascii files
    # reusable bits
    header = "ncols        720\nnrows        360\nxllcorner    -180\nyllcorner    -90\ncellsize     0.5\nNODATA_value -999\n"
    rowOrder = range(360)
    rowOrder.reverse()
    
    # get a set of dates to put in filename
    if average:
        if datefile == '':
            dateKeys = ['ave-' + str(dates[0][0]) + '.' + str(dates[0][1]) + '-' + str(dates[-1][0]) + '.' + str(dates[-1][1]) ]
        else:
            dateKeys = ['ave-' + os.path.basename(datefile)]
    else:
        # replace commas in dateKeys by '.'
        dateKeys = [str(m[0]) + '.' + str(m[1]) for m in dates]
    
    # get the grid indices from gridindex
    gridfile= transposedb + '/' + 'transposedGridRefIndex.txt'
    try:
        gridind = open(gridfile, 'r')
    except IOError:                                                  
        sys.stderr.write('Cannot open grid index file: ' + gridindex + '\n')
        sys.exit()
    
    gridindex = gridind.readlines()
    gridind.close()
    gridindex = [ x.rstrip() for x in gridindex]
    gridindex = [ x.split(',') for x in gridindex]
    gridindex = [ [int(row), int(col)]  for row, col in gridindex]
    
    if verbose: sys.stdout.write('Writing output files for '+ str(len(dateKeys)) + ' ascii files\n')
    
    # iterate over dates    
    for eachkey in range(len(dateKeys)):
        
        # generate a grid prefilled with -999
        ascii = numpy.zeros([360,720]) - 999
        
        # fill in asciiList
        for c in range(len(gridindex)):
            row = gridindex[c][1] - 1
            col = gridindex[c][0] - 1
            ascii[row , col] = valArray[eachkey, c]
        
        # write out asciiList as a grid
        outfile = outbase + dateKeys[eachkey] + '.asc'
        
        try:
            out = open(outfile, 'w')
        except IOError:                                                  
            sys.stderr.write('Cannot open output file: ' + outfile + '\n')
            sys.exit()
        
        # write header to file
        out.write(header)
        
        # write rows to file
        for r in rowOrder:
            # space delimited strings
            rowVals = [' '+ str(round(v, 2)) for v in ascii[r,]]
            rowVals = ''.join(rowVals)
            # lose first space 
            rowVals = rowVals[1:]
            out.write(rowVals + '\n')
        
        out.close()
        
        # file ticker
        if verbose: 
            sys.stdout.write('.')
            sys.stdout.flush()
    
    # end ticker
    sys.stdout.write('\n')


def usage():
    
    print """
  
  tcclimCruTS2:
  A data extraction tool for CRU TS 2.10 global climate data.
  
  Basic use:
     tcclimCruTS2.py [options] 'variable' 'startdate' 'enddate'
  
  Options:
     -h or --help:        show this help information
     -q or --quiet:       suppress normal progress reporting
     -m or --monthly:     provide monthly grid files not a single average file
     -d or --datefile:    use dates from provided datefile
     -o or --outbase:     use provided base name for output files
     -t or --transposedb: use provided path to transposed data folder 
  
  Available variables:
     'dtr','pre','tmp','vap','cld','tmn','tmx','frs','wet'
  Date format and range:
     '01,1901' to '12,2002'
  
  Returns ArcGIS ascii grid files of global data for a given variable. The 
  user can provide a range of dates from the command line. Alternatively, 
  the user can name a file containing dates, with each line holding a single
  date in the format above. Dates provided in this way need not be consecutive.
  
  By default, the program calculates an average across the dates selected
  but monthly files may be requested instead. As each output file is around
  2 MB, this can be slow and use a lot of disk space for large date ranges.
  
  The tool requires a database of transposed grim files created using the 
  program 'tcclimDownloadCRUTS2.py'. This transposition radically improves 
  handling speed for converting the data to ascii grids. By default, this 
  program looks for a folder called 'tcclim_cru_ts_2_10' in the directory 
  from which the program is being run.
  
  Example use:
     # Get precipitation from December 2002
     # This is the confirmed datum for checking purposes from:
     # http://www.cru.uea.ac.uk/~timm/grid/CRU_TS_2_1.html 
     # Grid cell: row 286, col 363 from lower left origin = 90.0
     tcclimCruTS2.py -m pre  12,2002  12,2002
    
     # Get mean precipitation from HadCM3 A2 from July 1971 to July 1991
     tcclimCruTS2.py pre  07,1971  07,1991
     
     # Get monthly grid files for precipitation for dates in a file
     tcclimCruTS2.py -m -d exampleDatesCruTS2.txt pre
    
    """


def main():
    
    # Tidies up and passes options to this function:
    # def getTynSC2(model, scenario, variable, start='', end='', 
    #               outbase='', datefile='', average=True, verbose=True, 
    #               transposedb='transpose'):
    
    # complex input to test control
    # python getTynSC2.py -mq -d dates.txt -o 'myFiles' --transposedb '../transpose' 'hadcm3' 'a2'  'pre'
    
    # sort out the options and argument
    try:
        opts, args = getopt.getopt(sys.argv[1:], "hqmd:o:t:", ["help", "quiet","monthly","datefile=","outbase=", "transposedb="])
    except getopt.GetoptError, err:
        # print help information and exit:
        print str(err) # will print something like "option -a not recognized"
        usage()
        sys.exit()
    
    # default options
    verbose=True
    average=True
    datefile=''
    outbase =''
    transposedb='tcclim_cru_ts_2_10'
    
    # initialise blank dates for handling datefile option
    start=''
    end=''
    
    # initial setups
    for o, a in opts:
        if o in ('-h', '--help'):
            usage()
            sys.exit()
        elif o in ('-q', '--quiet'):
            verbose=False
        elif o in ('-m', '--monthly'):
            average=False
        # getopt handles missing arguments for d, o and t but can easily pass model
        # and scenario names in as datefiles. Let getTynSC2 function handle this
        elif o in ('-d', '--datefile'):
            datefile = a
        elif o in ('-o', '--outbase'):
            outbase = a
        elif o in ('-t', '--transposedb'):
            transposedb = a
        else:
            assert False, "Unhandled option"
            usage()
            sys.exit()
    
    if len(args) == 1:
        var = args[0]
    elif len(args) == 3:
        var = args[0]
        start = args[1]
        end = args[2]
    else:
        sys.stderr.write('\n  Incorrect number of arguments:\n')
        sys.stderr.write('  tcclimCruTS2 require either 1 argument (variable) and a datefile\n')
        sys.stderr.write('  or three arguments (variable, start date, end date)\n')
        usage()
        sys.exit()
    
    ## [sys.stdout.write(o + ',' + a + '\n') for o, a in opts]
    ## [sys.stdout.write(a + '\n') for a in args]
    
    getCruTS2(variable=var, start=start, end=end,  
              outbase=outbase, datefile=datefile, average=average, verbose=verbose,
              transposedb=transposedb)


if __name__ == "__main__":
    main()
